# name_generator.py
import asyncio
import json
import time
from config import Config
from db import DB
from ai_client import AIClientAsync
import logging
from tqdm import tqdm
import datetime


# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 系统消息 (包含所有静态规则)
SYSTEM_MESSAGE = """
你是名字生成专家，会生成绝不重复的名字，并且严格遵守名字长度要求。

请生成一个中文名字的数组，要求：
1. 核心原则：
   - 必须是完整的中文名字（包括姓氏），不能是单个的字。
   - 必须完全符合中文姓名美学和规范
   - 每个名字需有美好寓意（吉祥、聪慧、高雅、独特等）
   - 绝对避免使用生僻字、歧义字、以及任何可能引起不适的字词,除非有必要，否则文字笔画不用太复杂，尽量使用较为简单的文字。
   - 名字不能重复
   - 名字音韵需要符合姓氏的音韵，读起来富有节奏感、易于记忆和传播

2. 具体规范：
   【禁用列表】
   × 负面含义字：如 凶、恶、伤、亡、死 等
   × 与粗俗、低俗相关的字
   × 过度常见的名字（避免过于流行的名字）
"""

# 提取 prompt 模板 (精简版)
NAME_GENERATION_PROMPT_TEMPLATE = """
请生成一个 {surname} 姓{gender}性名字的数组，要求：

1. 姓氏：**{surname}**
2. 性别: **{gender}**
3. 姓名长度: **包含姓氏总共{name_length}个字, 必须严格遵守长度限制。绝对不能生成{name_length_error}个字的名字!**
4. 风格要求：**{style_description}**
5. 数量: {batch_size}个
6. 格式要求：
   - **输出必须是纯粹的、唯一的 JSON 数组格式。除了数组本身，绝对不能包含任何其他字符、单词、解释、说明或文本，任何不符合JSON格式的内容都将被视为无效。**
   - **正确的格式示例：** `["张三", "李四"]` (假设生成两个字的名字)
   - **请务必严格校验每个名字是否符合所有要求, 任何不符合的名字都应被排除**
"""

class NameGenerator:
    def __init__(self, config: Config, db: DB, ai_client: AIClientAsync):
        self.config = config
        self.db = db
        self.ai_client = ai_client
        self.styles = [  # 这里的示例也需要注意长度
  {
    "short": "传统",
    "long": "传统风格：注重寓意吉祥、和谐、稳重，用字规范典雅，体现文化底蕴和家族传承。名字整体音韵和谐，朗朗上口，避免生僻字和不雅谐音。"
  },
  {
    "short": "现代",
    "long": "现代风格：用字简洁、时尚、明快，具有创新性和个性，不拘泥于传统，可适当使用叠字或外文谐音，但需注意整体协调性和美感，避免过于标新立异或产生歧义。"
  },
  {
    "short": "自然",
    "long": "自然风格：名字与自然景物、元素巧妙关联，体现清新、恬静、自然、灵动之美。避免直接使用具体自然景物名称，而是采用与之相关的、寓意美好的字词，使名字整体意境优美、富有诗意。"
  },
  {
    "short": "古典",
    "long": "古典风格：名字取材于古代诗词、典籍、历史故事，用字古朴、典雅，具有浓厚的文化气息和历史感。名字整体需体现出古典韵味，避免生僻字和难以理解的典故。"
  },
  {
    "short": "神话",
    "long": "神话风格：名字可借鉴神话传说、古代神祇、神兽的特征或相关元素，营造神秘、奇幻、浪漫的氛围。避免直接使用神祇或神兽的全名，而是巧妙提取其中寓意美好的字词，或采用象征、比喻等手法，使名字既有文化内涵，又不失现代感。"
  },
 {
    "short": "动物",
    "long": "动物风格：以动物为灵感，通过隐喻、象征、谐音等方式，将动物的美好特质融入名字中，体现女性的柔美、灵动、聪慧。避免使用任何动物的名称。优先选用动物雅称，或提取动物相关特征中，具有美好寓意的字词。"
  },
  {
    "short": "色彩",
    "long": "色彩风格：名字中融入色彩相关的字词，以色彩的象征意义来体现不同的意境和气质。避免过于直白，应注重色彩与名字整体意境的融合。"
  },
  {
    "short": "数字",
    "long": "数字风格：巧妙地将数字融入名字中，使其具有独特的含义和趣味性。数字的运用应简洁明了，避免生硬的堆砌和过于复杂的解释。优先一位数字"
  },
  {
    "short": "国际",
    "long": "国际风格：名字中融入外文元素，可采用外文名的谐音，体现国际化、时尚感。注意中外文元素的协调性，避免生硬的组合和文化冲突。优先考虑常见的英文名。"
  },
  {
    "short": "诗意",
    "long": "诗意风格：名字富有诗情画意，用字优美、灵动、有意境，常引用或化用诗词典故，但需避免生僻字和过于晦涩的典故。名字整体应给人以美的享受，体现出文化底蕴。"
  }
]
        self.genders = ['男', '女']
        self.name_lengths = [2, 3]
        self.existing_names_cache = {}
        self.semaphore = asyncio.Semaphore(config.MAX_CONCURRENCY)

    def get_style_choices(self):
        return [style["short"] for style in self.styles]

    async def generate_names(self, reset=False, surname=None, length=None, style=None):
        start_time = time.time()
        logging.info("开始批量生成名字...")

        if reset:
            await self.db.execute("UPDATE surnames SET is_active = 1 WHERE is_active = 0")
            logging.info("已重置所有姓氏的 is_active 状态。")

        where_clauses = ["is_active = 1", "is_approved = 1"]
        params = []

        if surname:
            where_clauses.append("surname = %s")
            params.append(surname)

        where_clause = " AND ".join(where_clauses)
        query = f"SELECT * FROM surnames WHERE {where_clause} ORDER BY id"
        surnames = await self.db.fetch_all(query, tuple(params))

        if surname:
            total_surnames = len(surnames)
        else:
            total_surnames = (await self.db.fetch_one(f"SELECT COUNT(*) as count FROM surnames WHERE {where_clause}", tuple(params)))['count']

        processed_surnames = 0

        with tqdm(total=total_surnames, desc="处理姓氏") as pbar:
            for surname_record in surnames:
                logging.info(f"开始处理姓氏：{surname_record['surname']}")
                await self.generate_names_for_surname(surname_record, length, style)
                processed_surnames += 1
                pbar.update(1)
                logging.info(f"已处理姓氏：{surname_record['surname']} ({processed_surnames}/{total_surnames})")

        end_time = time.time()
        duration = end_time - start_time
        logging.info(f"批量生成名字完成！总耗时：{duration:.2f} 秒")


    async def generate_names_for_surname(self, surname, specified_length=None, specified_style=None):
        tasks = []
        for gender in self.genders:
            lengths = [specified_length] if specified_length else self.name_lengths
            for name_length in lengths:
                styles = [s for s in self.styles if s['short'] == specified_style] if specified_style else self.styles
                for style in styles:
                    existing_count = await self.check_existing_names_count(surname['surname'], gender, style['short'], name_length)
                    if existing_count >= 100:
                        logging.info(f"  - 跳过：性别={gender}，长度={name_length}，风格={style['short']} (已存在 {existing_count} 个，超过 100)")
                        continue

                    logging.info(f"  - 正在处理：性别={gender}，长度={name_length}，风格={style['short']} (已存在 {existing_count} 个)")
                    remaining_count = 100 - existing_count
                    num_batches = (remaining_count + self.config.BATCH_SIZE - 1) // self.config.BATCH_SIZE
                    for _ in range(num_batches):
                        task = asyncio.create_task(
                            self.generate_names_for_combination(surname['surname'], gender, style, name_length, min(self.config.BATCH_SIZE, remaining_count))
                        )
                        tasks.append(task)
                        remaining_count -= self.config.BATCH_SIZE

        await asyncio.gather(*tasks)
        await self.db.execute("UPDATE surnames SET is_active = 0 WHERE id = %s", (surname['id'],))
        self.clear_cache()



    async def check_existing_names_count(self, surname, gender, style, name_length):
        cache_key = f"{surname}_{gender}_{style}_{name_length}"
        if cache_key in self.existing_names_cache:
            return self.existing_names_cache[cache_key]

        count = (await self.db.fetch_one(
            "SELECT COUNT(*) as count FROM names WHERE surname = %s AND gender = %s AND style = %s AND name_length = %s",
            (surname, gender, style, name_length)
        ))['count']
        self.existing_names_cache[cache_key] = count
        return count

    def clear_cache(self):
        self.existing_names_cache = {}

    async def generate_names_for_combination(self, surname, gender, style, name_length, count):
        # 针对不同长度，给出不同的错误长度提示
        name_length_error = 3 if name_length == 2 else 2

        style_description = style['long']
        prompt = NAME_GENERATION_PROMPT_TEMPLATE.format(
            surname=surname,
            gender=gender,
            style_description=style_description,
            name_length=name_length,
            name_length_error=name_length_error,  # 动态错误长度
            batch_size=count
        )

        messages = [
            {'role': 'system', 'content': SYSTEM_MESSAGE},
            {'role': 'user', 'content': prompt},
        ]

        async with self.semaphore:
            retry_count = 0
            while retry_count < self.config.MAX_RETRIES:
                logging.info(f"    - 准备生成名字：剩余需生成数量={count}，重试次数={retry_count}，长度={name_length}")
                try:
                    names = await self.ai_client.chat(messages, name_length=name_length)

                    if names is None:
                        logging.warning("    - ⚠️ AI 返回 None，可能是模型不支持此长度。")
                        return

                    if not names:
                        logging.warning("    - ⚠️ AI 没有返回任何名字。")
                        retry_count += 1
                        messages.append({'role': 'user', 'content': '请生成一批全新的名字，并且严格遵守名字长度要求。'})
                        continue

                    valid_names = []
                    if isinstance(names, list):
                        for name in names:
                            if len(name) == name_length:
                                valid_names.append(name)
                            else:
                                logging.info(f"    - ❌ 名字“{name}”长度不符合要求（应为 {name_length}），已忽略")
                    else:
                         logging.warning(f"    - ⚠️ AI 返回的不是列表: {names}")
                         retry_count += 1
                         continue

                    placeholders = ', '.join(['%s'] * len(valid_names))
                    query = f"SELECT name FROM names WHERE surname = %s AND name IN ({placeholders})"
                    existing_names_in_db = await self.db.fetch_all(query, (surname, *valid_names))
                    existing_names_in_db = [item['name'] for item in existing_names_in_db]

                    new_names = [name for name in valid_names if name not in existing_names_in_db]
                    new_names = list(set(new_names))

                    logging.info(f"    - AI 返回名字数：{len(names)}，有效名字数：{len(valid_names)}，去重后：{len(new_names)}")

                    if new_names:
                        now = datetime.datetime.now()
                        values = [
                            (name, surname, gender, style['short'], name_length, None, now)
                            for name in new_names
                        ]

                        await self.db.insert_many(
                            "INSERT IGNORE INTO names (name, surname, gender, style, name_length, explanation, created_at) VALUES (%s, %s, %s, %s, %s, %s, %s)",
                                values
                        )
                        logging.info(f"    - ✅ 已成功插入 {len(new_names)} 个名字。")
                        return

                    else:
                        logging.warning("    - ⚠️ AI 返回的名字全部无效或已存在。")
                        retry_count+=1
                        messages.append({'role': 'user', 'content': '请生成一批全新的名字，并且严格遵守名字长度要求。'})

                except Exception as e:
                    logging.error(f"    - ❌ 生成名字时出错（重试 {retry_count}）：{e}")
                    retry_count += 1
                    messages.append({'role': 'user', 'content': '请生成一批全新的名字，并且严格遵守名字长度要求。'})
            logging.error(f"    - ❌ 经过 {self.config.MAX_RETRIES} 次重试，仍未能生成足够数量的有效名字。")